###### Bayesian IRT

# Following McGann (2014)

library(R2jags)
library(tidyverse)
library(coda)

Data <- readRDS(file = "/Users/matteo/Desktop/Prison Polling/AttitudesToCrimeAndPunishment.RData")

Data[[1]] <- Data[[1]][Data[[1]]$Demographic == "All adults",]
Data[[2]] <- Data[[2]][Data[[2]]$Demographic == "All adults",]
Data[[3]] <- Data[[3]][Data[[3]]$Demographic == "All adults",]
Data[[4]] <- Data[[4]][Data[[4]]$Demographic == "All adults",]
Data[[4]] %>% mutate(n = 1001) -> Data[[4]] # assumed constant n for all survey questions

Do_IRT <- function (x) {
  Data[[x]] %>% 
    mutate(year = as.numeric(substr(as.character(.$Date),1,4))) %>%
    group_by(year, Varname) %>%
    summarise(Index = weighted.mean(Index, n),
              n = sum(n)) %>%
    ungroup() -> data
  
  #Convert question names from R factors to numeric
  data$q <- as.numeric(as.factor(data$Varname))
  
  #Create list for JAGS
  dataList <- c(as.data.frame(data), list(len=nrow(data),
                                          nquest=length(unique(data$Varname)),
                                          yearlist=as.numeric(unique(data$year))))
  
  jags.params = c("mu","sigma",
                  "lambda", "alpha", "b")
  
  jagsfit <- jags(data = dataList, 
                  inits = NULL, 
                  parameters.to.save = jags.params, 
                  n.iter = 100000, 
                  model.file = "/Users/matteo/Desktop/Prison Polling/CrimeModel_JAGS.txt", 
                  n.chains = 3)
  
  return(jagsfit)
}

fitted <- lapply(c(1,2,3,4), Do_IRT)



#### convergence check

title <- c('fear','punitiveness','death penalty','prioritisation')

as_tibble(gelman.diag(as.mcmc(fitted[[1]]))$psrf, rownames = 'varname') |>
  filter(str_detect(varname, c("^alpha|^lambda|^sigma|^mu|b"))) |>
  mutate(group = case_when(str_detect(varname, c("^alpha")) ~ 'alpha',
                           str_detect(varname, c("^lambda")) ~ 'lambda',
                           str_detect(varname, c("^mu")) ~ 'mu',
                           str_detect(varname, c("^sigma")) ~ 'sigma',
                           str_detect(varname, c("^b")) ~ 'b')) |>
  pivot_longer(!c(varname,group)) |>
  ggplot(aes(x = value, colour = name, y = varname, group = varname)) +
  geom_line(colour = 'black') + geom_point() + geom_vline(xintercept = 1.1) +
  xlab('potential scale reduction factor') + ylab(NULL) +
  facet_wrap(facets = vars(group), scales = "free_y", ncol = 2) +
  labs(title = title[[1]])

as_tibble(gelman.diag(as.mcmc(fitted[[1]]))$psrf, rownames = 'varname') |>
  filter(str_detect(varname, c("^alpha|^lambda|^sigma|^mu|b"))) |>
  mutate(group = case_when(str_detect(varname, c("^alpha")) ~ 'alpha',
                           str_detect(varname, c("^lambda")) ~ 'lambda',
                           str_detect(varname, c("^mu")) ~ 'mu',
                           str_detect(varname, c("^sigma")) ~ 'sigma',
                           str_detect(varname, c("^b")) ~ 'b')) |>
  arrange(desc(`Point est.`))

stack(effectiveSize(as.mcmc(fitted[[2]]))) |>
  filter(str_detect(ind, c("^alpha|^lambda|^sigma|^mu|b"))) |>
  mutate(group = case_when(str_detect(ind, c("^alpha")) ~ 'alpha',
                           str_detect(ind, c("^lambda")) ~ 'lambda',
                           str_detect(ind, c("^mu")) ~ 'mu',
                           str_detect(ind, c("^sigma")) ~ 'sigma',
                           str_detect(ind, c("^b")) ~ 'b')) |>
  pivot_longer(!c(ind,group)) |>
  ggplot(aes(x = value, y = ind)) +
  geom_point() + xlab('effective sample size') + ylab(NULL) + 
  geom_vline(xintercept = 1000) +
  facet_wrap(facets = vars(group), scales = "free_y", ncol = 2) +
  labs(title = title[[2]])

# All four models have converged



#### Appendix 3: IRT estimates plot

title <- c("IRT estimates of Crime Concern","IRT estimates of Punitiveness",
           "IRT estimates of Death Penalty","IRT estimates of Prioritisation")

MCMCvis::MCMCsummary(fitted[[1]], 
                     params = 'mu', 
                     Rhat = FALSE, 
                     n.eff = FALSE, 
                     HPD = TRUE, 
                     hpd_prob = 0.89) |>
  rownames_to_column('varname') |>
  mutate(Year = as.numeric(substr(varname, 4,7))) |>
  ggplot(aes(x = Year, group = Year)) +
  geom_errorbar(aes(ymin = `89%_HPDL`, ymax = `89%_HPDU`)) + 
  geom_point(aes(y = mean)) +
  xlab(NULL) + ylab(NULL) + scale_x_continuous(limits = c(1965,2023)) +
  labs(title = title[[1]]) + 
  MyThemes::theme_base() + theme(plot.title = element_text(face = 'bold', hjust = 0.5)) -> p1
MCMCvis::MCMCsummary(fitted[[2]], 
                     params = 'mu', 
                     Rhat = FALSE, 
                     n.eff = FALSE, 
                     HPD = TRUE, 
                     hpd_prob = 0.89) |>
  rownames_to_column('varname') |>
  mutate(Year = as.numeric(substr(varname, 4,7))) |>
  ggplot(aes(x = Year, group = Year)) +
  geom_errorbar(aes(ymin = `89%_HPDL`, ymax = `89%_HPDU`)) + 
  geom_point(aes(y = mean)) +
  xlab(NULL) + ylab(NULL) + scale_x_continuous(limits = c(1981,2023)) +
  labs(title = title[[2]]) + 
  MyThemes::theme_base() + theme(plot.title = element_text(face = 'bold', hjust = 0.5)) -> p2
MCMCvis::MCMCsummary(fitted[[3]], 
                     params = 'mu', 
                     Rhat = FALSE, 
                     n.eff = FALSE, 
                     HPD = TRUE, 
                     hpd_prob = 0.89) |>
  rownames_to_column('varname') |>
  mutate(Year = as.numeric(substr(varname, 4,7))) |>
  ggplot(aes(x = Year, group = Year)) +
  geom_errorbar(aes(ymin = `89%_HPDL`, ymax = `89%_HPDU`)) + 
  geom_point(aes(y = mean)) +
  xlab(NULL) + ylab(NULL) + scale_x_continuous(limits = c(1962,2023)) +
  labs(title = title[[3]]) + 
  MyThemes::theme_base() + theme(plot.title = element_text(face = 'bold', hjust = 0.5)) -> p3
MCMCvis::MCMCsummary(fitted[[4]], 
                     params = 'mu', 
                     Rhat = FALSE, 
                     n.eff = FALSE, 
                     HPD = TRUE, 
                     hpd_prob = 0.89) |>
  rownames_to_column('varname') |>
  mutate(Year = as.numeric(substr(varname, 4,7))) |>
  ggplot(aes(x = Year, group = Year)) +
  geom_errorbar(aes(ymin = `89%_HPDL`, ymax = `89%_HPDU`)) + 
  geom_point(aes(y = mean)) +
  xlab(NULL) + ylab(NULL) + scale_x_continuous(limits = c(1973,2023)) +
  labs(title = title[[4]]) + 
  MyThemes::theme_base() + theme(plot.title = element_text(face = 'bold', hjust = 0.5)) -> p4

cowplot::plot_grid(p1,p2,p3,p4, nrow = 2)

saveRDS(fitted, file = "/Users/matteo/Desktop/Prison Polling/IRTModels.RData")
